import numpy as np
import random

# 将单个数据点分配到最近的聚类中心
def assign_cluster(x, centroids):
    # 计算数据点到每个聚类中心的欧氏距离
    distances = np.sqrt(np.sum((centroids - x) ** 2, axis=1))
    return np.argmin(distances)

# K均值聚类算法的主函数
def Kmeans(data, k, epsilon=1e-6, max_iteration=100):

    data = np.array(data)
    n_samples, n_features = data.shape

    indices = random.sample(range(n_samples), k)
    centroids = data[indices]

    labels = np.zeros(n_samples)

    # 迭代
    for i in range(max_iteration):
        # 保存旧的聚类中心
        old_centroids = centroids.copy()

        # 分配每个数据点到最近的聚类中心
        for j in range(n_samples):
            labels[j] = assign_cluster(data[j], centroids)

        # 更新聚类中心
        for c in range(k):
            # 找到属于当前聚类的所有数据点
            cluster_points = data[labels == c]
            # 如果聚类中没有数据点，则重新随机选择一个数据点作为中心
            if len(cluster_points) == 0:
                centroids[c] = data[random.randint(0, n_samples - 1)]
            else:
                # 计算聚类中心的平均值
                centroids[c] = np.mean(cluster_points, axis=0)

        # 检查收敛条件
        centroid_shift = np.sum(np.sqrt(np.sum((centroids - old_centroids) ** 2, axis=1)))
        if centroid_shift < epsilon:
            print(f"算法在第{i + 1}次迭代收敛")
            break

    if i == max_iteration - 1:
        print(f"算法达到最大迭代次数{max_iteration}")

    return centroids, labels